#! /usr/bin/env python3
import os.path
import cv2
import apriltag
import cv2.aruco

def video_pic_det(path):
    if path.endswith('.jpg') or path.endswith('.png'):
        frame=cv2.imread(path)
        image = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
        detector = apriltag.Detector(apriltag.DetectorOptions(families="tag36h11"))

        detections = detector.detect(image)
        print(detections)
        for i in detections:
            x,y=i.center
            tag_id=i.tag_id
            cv2.line(frame,(round(i.corners[0][0]),round(i.corners[0][1])),(round(i.corners[1][0]),round(i.corners[1][1])),(0,255,0),1)
            cv2.line(frame, (round(i.corners[0][0]), round(i.corners[0][1])),
                     (round(i.corners[3][0]), round(i.corners[3][1])), (0, 255, 0), 1)
            cv2.line(frame, (round(i.corners[1][0]), round(i.corners[1][1])),
                     (round(i.corners[2][0]), round(i.corners[2][1])), (0, 255, 0), 1)
            cv2.line(frame, (round(i.corners[2][0]), round(i.corners[2][1])),
                     (round(i.corners[3][0]), round(i.corners[3][1])), (0, 255, 0), 1)

        cv2.namedWindow('ori', cv2.WINDOW_NORMAL)  # 窗口大小可以改变
        cv2.namedWindow('dst', cv2.WINDOW_NORMAL)
        cv2.imshow('ori',frame)
        cv2.imshow('dst',image)
        key=cv2.waitKey(0)
        if key & 0xFF == ord('q'):  # ‘q’键实现退出
            cv2.destroyAllWindows()
            exit(0)

    else:
        # 打开摄像头
        cap = cv2.VideoCapture(path)

        while (cap.isOpened()):
            # 开始用摄像头读数据，返回hx为true则表示读成功，frame为读的图像
            ret,frame = cap.read()

            if ret:
                image = cv2.cvtColor(frame,cv2.COLOR_RGB2GRAY)
                detector = apriltag.Detector(apriltag.DetectorOptions(families="tag36h11"))
                detections = detector.detect(image)
                for i in detections:
                    x, y = i.center
                    tag_id = i.tag_id
                    cv2.line(frame, (round(i.corners[0][0]), round(i.corners[0][1])),
                             (round(i.corners[1][0]), round(i.corners[1][1])), (0, 255, 0), 1)
                    cv2.line(frame, (round(i.corners[0][0]), round(i.corners[0][1])),
                             (round(i.corners[3][0]), round(i.corners[3][1])), (0, 255, 0), 1)
                    cv2.line(frame, (round(i.corners[1][0]), round(i.corners[1][1])),
                             (round(i.corners[2][0]), round(i.corners[2][1])), (0, 255, 0), 1)
                    cv2.line(frame, (round(i.corners[2][0]), round(i.corners[2][1])),
                             (round(i.corners[3][0]), round(i.corners[3][1])), (0, 255, 0), 1)
                # print(detections)
                cv2.namedWindow('ori', cv2.WINDOW_NORMAL)  # 窗口大小可以改变
                cv2.namedWindow('dst', cv2.WINDOW_NORMAL)
                cv2.imshow('ori', frame)
                cv2.imshow('dst', image)
                key = cv2.waitKey(1)
                if key & 0xFF == ord('q'):  # ‘q’键实现退出
                    cv2.destroyAllWindows()
                    exit(0)

def distance_(p1:(int,int),p2:(int,int)):
    fx = 748.454
    fy = 754.842
    cx = 323.751
    cy = 207.372
    x1,y1=p1
    x2,y2=p2
    u1, v1 = (x1 - cx) / fx, (y1 - cy) / fy
    u2, v2 = (x2 - cx) / fx, (y2 - cy) / fy
    d = ((u2 - u1) ** 2 + (v2 - v1) ** 2) ** 0.5
    l = 28.7 - 4.04 - 4.04#实际长度 mm
    z = l / d
    return z

def aruco_detect(cap_ind):
    cap = cv2.VideoCapture(cap_ind)
    conor_dict={}

    def real_mm(markerCorners, markerIds):
        for ind, val in enumerate(markerCorners):
            r_mm=57.4
            # total_id_num:24 0-23
            conor_dict[ind] = val[0][0]


            # print(f'{markerIds[ind][0]}: , {val}')
        print(conor_dict)
        for key in conor_dict.keys():
            if key%2==0:
                key2=key+1
                if key2 in conor_dict:
                    a2=(conor_dict[key2][0]-conor_dict[key][0])**2
                    b2=(conor_dict[key2][1]-conor_dict[key][1])**2
                    c=(a2+b2)**0.5
                    pix_mm=28.7*2/c


    while (cap.isOpened()):
        # 开始用摄像头读数据，返回hx为true则表示读成功，frame为读的图像
        ret, frame = cap.read()

        if ret:
            image = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
            dictionary = cv2.aruco.getPredefinedDictionary(cv2.aruco.DICT_4X4_50)
            parameters = cv2.aruco.DetectorParameters()
            detector = cv2.aruco.ArucoDetector(dictionary, parameters)



            markerCorners, markerIds, rejectedCandidates = detector.detectMarkers(image)
            real_mm(markerCorners, markerIds)
            # 画出标志位置
            cv2.aruco.drawDetectedMarkers(frame, markerCorners, markerIds)

            cv2.namedWindow('ori', cv2.WINDOW_NORMAL)  # 窗口大小可以改变
            cv2.namedWindow('dst', cv2.WINDOW_NORMAL)
            cv2.imshow('ori', frame)
            cv2.imshow('dst', image)
            key = cv2.waitKey(1)
            if key & 0xFF == ord('q'):  # ‘q’键实现退出
                cv2.destroyAllWindows()
                exit(0)




if __name__ == "__main__":
    video_dir = "/home/yuan/Videos/video"
    pic_dir = '/home/yuan/Pictures'
    # name = '1920-1080-mjpg-60.png'
    # name="1920-1080-60fps-mjpg.avi"
    # name="640-480or360-60fps-mjpg.avi"
    # name="640-480-mjpg-60-2-part.png"
    # name="1280-720-yuy2-fps10.avi"
    # name="new-1920-1080-mjpg.avi"

    # name="n_640-mjpg.avi"
    # name="n_1080*720-mjpg.avi"
    video_dir='/home/yuan/xl_ws/src/grabber_ws'
    name='output_video.mp4'
    path = os.path.join(video_dir, name)
    video_pic_det(path)
    # aruco_detect(0)



